from utils.basic_function import *
from utils.database_table import FusedRoadTable, RoadTable
from geopandas.geoseries import *


class TopologyTool:
    def __init__(self, node_table, tolerance):
        self.node_table = node_table
        self.tolerance = tolerance
        self.scan_order = []
        self.road_dict = dict()

    def label_nodes_workflow(self, extent=None):
        for r_class in self.scan_order:
            road_table = self.road_dict[r_class]
            road_table.add_column(road_table.snodeid, set_null=True)
            road_table.add_column(road_table.enodeid, set_null=True)
            self.label_nodes(r_class, extent)

    def label_nodes(self, r_class, extent):
        """
        遍历遥感路网table中的所有路段，将node插入node表。如果端点容差范围内已有node，则沿用此node。

        Args:
            r_class: 道路等级
            extent: Polygon,构建拓扑的范围

        Returns:

        """
        road_table = self.road_dict[r_class]
        if extent:
            where = f"""
            where st_intersects({road_table.geom}, 
            st_geomfromewkt('SRID={road_table.srid};{extent.wkt}')
            )
        """
        else:
            where = ""

        sql = f"select {road_table.pk}, {road_table.geom} from {road_table.name} {where} order by {road_table.pk} asc"
        gdf = geopd.read_postgis(sql, conn)

        # 遍历路段
        for idx in gdf.index:
            geom = gdf.loc[idx, road_table.geom]
            pk = gdf.loc[idx, road_table.pk]
            st_coord = geom.coords[0]
            ed_coord = geom.coords[-1]
            snodeid = self.label_node(st_coord, r_class, pk)
            enodeid = self.label_node(ed_coord, r_class, pk)
            # 更新路段表里snodeid、enodeid属性
            road_table.update_attr(road_table.pk, pk, {road_table.snodeid: snodeid, road_table.enodeid: enodeid})
            print(f"{r_class}, {pk}")

    def label_node(self, coord, current_class, pk_val):
        pass

    @staticmethod
    def point_interpolate(linestring, pt):
        length = linestring.length
        line_coord_list = list(linestring.coords)
        before_length = 0
        proj_result = project_pt_to_polyline(pt.coords[0], line_coord_list)
        for i in range(proj_result["pre_node"]):
            before_length += cal_dist(line_coord_list[i], line_coord_list[i + 1])
        before_length += cal_dist(line_coord_list[proj_result["pre_node"]], proj_result["proj_pt"])
        return before_length / length

    def get_mid_intersection(self, coord, gdf, geom_col):
        """
        计算coord到gdf中每条路段端点的距离和最短距离。若coord不在路段的端点上（在路段中间），返回相应的路段。

        Args:
            coord:
            gdf:
            geom_col:

        Returns:

        """
        if gdf.empty:
            return gdf
        gdf["st_dist"] = gdf[geom_col].apply(lambda x: Point(x.coords[0]).distance(Point(coord[0], coord[1])))
        gdf["ed_dist"] = gdf[geom_col].apply(lambda x: Point(x.coords[-1]).distance(Point(coord[0], coord[1])))
        gdf["mid_dist"] = gdf[geom_col].apply(lambda x: x.distance(Point(coord[0], coord[1])))
        gdf["dist"] = gdf.apply(lambda x: min([x["st_dist"], x["ed_dist"], x["mid_dist"]]), axis=1)
        gdf.loc[gdf["st_dist"] < self.tolerance, "intersect_type"] = "st"
        gdf.loc[gdf["ed_dist"] < self.tolerance, "intersect_type"] = "ed"
        gdf["intersect_type"] = gdf["intersect_type"].fillna("mid")
        intersect_gdf = gdf[(gdf["intersect_type"] == "mid") & (gdf["mid_dist"] < self.tolerance)].sort_values("dist", ascending=True)
        return intersect_gdf

    def get_roads_by_node(self, node_id):
        """
        获取端点为某结点的所有路段

        Args:
            node_id:结点编号

        Returns:
            independent_df: gid：路段编号, class：路段类型, is_st：bool，该结点是否是起点

        """
        # 端点为该结点的路段
        gdf_list = []
        for r_class in self.scan_order:
            road_table = self.road_dict[r_class]
            temp_gdf = pd.read_sql(f"""
                select {road_table.pk}, {road_table.snodeid}, {road_table.enodeid} from {road_table.name} 
                where {road_table.snodeid} = {node_id} or {road_table.enodeid} = {node_id}""", conn)
            if not temp_gdf.empty:
                temp_gdf["class"] = r_class
                gdf_list.append(temp_gdf)
        independent_df = pd.concat(gdf_list).reset_index(drop=True)
        independent_df["is_st"] = False
        independent_df.loc[independent_df[road_table.snodeid] == node_id, "is_st"] = True
        independent_df.drop(columns=[road_table.snodeid, road_table.enodeid], inplace=True)
        return independent_df

    def get_superior_roads_by_node(self, node_id):
        """
        获取该结点依赖的路段

        Returns:
            dependent_df: gid：路段编号, class：路段类型
        """
        # 该结点依赖的路段
        dependent_df = pd.read_sql(f"""
                    select {self.node_table.superior_pk}, {self.node_table.superior_class} 
                    from {self.node_table.name} 
                    where {self.node_table.pk} = {node_id}
                """, conn)
        if dependent_df.loc[0, self.node_table.superior_pk] is not None:
            dependent_pk = [int(x) for x in dependent_df.loc[0, self.node_table.superior_pk].split("|")]
            dependent_class = dependent_df.loc[0, self.node_table.superior_class].split("|")
        else:
            dependent_pk = []
            dependent_class = []

        dependent_df = pd.DataFrame({RoadTable.pk: dependent_pk, "class": dependent_class})
        return dependent_df

    def get_nodes_by_road(self, r_class, road_pk):  # 包括端点和依赖点
        nodes_dict = {}
        # 查询road_pk对应的道路，获取端点
        road_table = self.road_dict[r_class]
        df = pd.read_sql(f"select {road_table.snodeid}, {road_table.enodeid} from {road_table.name} where {road_table.pk} = {road_pk}", conn)
        if df.empty:
            return None
        nodes_dict[road_table.snodeid] = df.loc[0, road_table.snodeid]
        nodes_dict[road_table.enodeid] = df.loc[0, road_table.enodeid]

        # 查询依赖于该道路的结点
        mid_nodes = []
        # 只依赖一条道路
        df_single_dependent = pd.read_sql(f"""
            select {self.node_table.pk} from {self.node_table.name} where {self.node_table.superior_pk} = '{road_pk}' """, conn)
        mid_nodes.extend(df_single_dependent[self.node_table.pk].values.tolist())
        # 依赖不止一条道路
        df_multi_dependent = pd.read_sql(f"""
            select {self.node_table.pk}, {self.node_table.superior_pk}, {self.node_table.superior_class} 
            from {self.node_table.name} where
            {self.node_table.superior_pk} like '%{road_pk}%' and {self.node_table.dependent_num}>1""", conn)
        # 展开
        for idx in df_multi_dependent.index:
            pk_list = [int(x) for x in df_multi_dependent.loc[idx, self.node_table.superior_pk].split("|")]
            class_list = df_multi_dependent.loc[idx, self.node_table.superior_class].split("|")
            temp = pd.DataFrame({self.node_table.superior_pk: pk_list, self.node_table.superior_class: class_list})
            temp = temp.query(f"{self.node_table.superior_pk}=={road_pk} and {self.node_table.superior_class}=='{r_class}'")
            if not temp.empty:
                mid_nodes.append(df_multi_dependent.loc[idx, self.node_table.pk])
        nodes_dict["midnode"] = mid_nodes
        return nodes_dict

    def get_related_roads_by_road(self, r_class, road_pk):
        """
        给定road_pk，查询所有与此路段关联的路段（共享端点或存在依赖）

        Returns:
            inferior_df:依赖于此路段的路段。gid:路段编号, class:路段类别, is_st:依赖结点是否是该路段的起点, nodeid:结点编号
            superior_df:此路段依赖的路段。gid:路段编号, class:路段类别, superior_st:依赖结点是否是road_pk的起点, node_id:结点编号
            share_df:与此路段共享端点的路段。gid:路段编号, class:路段类别, is_st:共享端点是否是该路段的起点, share_st:共享结点是否是road_pk的起点, node_id:结点编号

        """
        # 获取该路段相关的结点（端点、中间点）
        nodes_dict = self.get_nodes_by_road(r_class, road_pk)
        if nodes_dict is None:
            return None, None, None

        # 依赖于此路段的路段
        inferior_list = []
        for node_id in nodes_dict['midnode']:  # 遍历依赖于此道路的结点
            temp_df = self.get_roads_by_node(node_id)  # 获取端点为这些结点的道路
            temp_df['nodeid'] = node_id
            inferior_list.append(temp_df)
        inferior_df = pd.concat(inferior_list).reset_index(drop=True)

        # 此路段依赖的路段
        snodeid = nodes_dict['snodeid']
        enodeid = nodes_dict['enodeid']
        snode_superior_df = self.get_superior_roads_by_node(snodeid)
        snode_superior_df['superior_st'] = True
        snode_superior_df['node_id'] = snodeid
        enode_superior_df = self.get_superior_roads_by_node(enodeid)
        enode_superior_df['superior_st'] = False
        enode_superior_df['node_id'] = enodeid
        superior_df = pd.concat([snode_superior_df, enode_superior_df]).reset_index(drop=True)
        superior_df[RoadTable.pk] = superior_df[RoadTable.pk].astype('int')

        # 与此路段共享端点的路段
        snode_share_df = self.get_roads_by_node(snodeid)
        snode_share_df = snode_share_df[snode_share_df[RoadTable.pk] != road_pk]  # 排除自身
        snode_share_df['share_st'] = True
        snode_share_df['node_id'] = snodeid
        enode_share_df = self.get_roads_by_node(enodeid)
        enode_share_df = enode_share_df[enode_share_df[RoadTable.pk] != road_pk]  # 排除自身
        enode_share_df['share_st'] = False
        enode_share_df['node_id'] = enodeid
        share_df = pd.concat([snode_share_df, enode_share_df]).reset_index(drop=True)

        return inferior_df, superior_df, share_df


class RSTopologyTool(TopologyTool):
    def __init__(self, node_table, layer_prefix, tolerance=CLUSTER_TOLERANCE_R):
        super(RSTopologyTool, self).__init__(node_table, tolerance)
        self.scan_order = R_CLASS_LIST
        self.road_dict = {}
        for r_class in self.scan_order:
            self.road_dict[r_class] = FusedRoadTable(layer_prefix + "_" + r_class)

    def label_node(self, coord, current_class, pk_val):
        temp_node = Point(coord[0], coord[1])
        # 先在node table里检索，若没有则在其他图层里检索是否与中段相交
        existed_nodes = self.node_table.buffer_intersect_by_pt(coord, self.tolerance)

        if not existed_nodes.empty:  # node table里已有该结点，返回已有结点编号
            existed_nodes["dist"] = existed_nodes[self.node_table.geom].apply(lambda x: x.distance(temp_node))
            row = existed_nodes.sort_values(by="dist", ascending=True).iloc[0]
            new_node_id = row[self.node_table.pk]
            return new_node_id

        else:  # node table里没有该结点，需要插入新结点
            intersect_gdf_list = []
            # 遍历所有图层，找所有可能的依赖路段
            for r_class in self.scan_order:
                match_gdf = self.road_dict[r_class].buffer_intersect_by_pt(coord, self.tolerance)
                geom_col = self.road_dict[r_class].geom
                if r_class.lower() == current_class.lower():
                    match_gdf = match_gdf[match_gdf[self.node_table.pk] != pk_val]  # 排除自身
                temp_gdf = self.get_mid_intersection(coord, match_gdf, geom_col)  # 返回与coord相交于中间的路段
                if not temp_gdf.empty:
                    temp_gdf["road_type"] = r_class
                    intersect_gdf_list.append(temp_gdf)

            if len(intersect_gdf_list) != 0:  # 依赖于某条道路。插入依赖结点
                intersect_gdf = pd.concat(intersect_gdf_list).reset_index(drop=True)
                # 构造结点属性
                if intersect_gdf.shape[0] == 1:  # 只依赖于一条道路
                    highest_class = intersect_gdf.loc[0, "road_type"][0]
                else:  # 同时依赖于多条道路，记录最高的级别
                    superior_class_list = intersect_gdf["road_type"].str[0]
                    highest_class = superior_class_list[0]
                    for r_class in superior_class_list[1:]:
                        if compare_road_rank(r_class, highest_class) == 1:
                            highest_class = r_class
                attr_dict = {self.node_table.superior_class: highest_class}
            else:  # 目前结点不存在依赖
                attr_dict = dict()

            # 插入新结点
            attr_dict[self.node_table.geom] = temp_node
            new_node_id = self.node_table.insert(attr_dict)
            return new_node_id
